import os
import random

import cv2
import h5py
import numpy as np
import torch
from scipy import ndimage
from scipy.ndimage.interpolation import zoom
from torch.utils.data import Dataset


def random_rot_flip(image, label):
    # 随机旋转
    k = np.random.randint(0, 4)
    image = np.rot90(image, k)
    label = np.rot90(label, k)
    # 随机翻转
    axis = np.random.randint(0, 2)
    image = np.flip(image, axis=axis).copy()
    label = np.flip(label, axis=axis).copy()
    return image, label


def random_rotate(image, label):
    # 随机旋转（-20°至20°）
    angle = np.random.randint(-20, 20)
    image = ndimage.rotate(image, angle, order=0, reshape=False)
    label = ndimage.rotate(label, angle, order=0, reshape=False)
    return image, label


class RandomGenerator(object):
    def __init__(self, output_size):
        self.output_size = output_size

    def __call__(self, sample):
        image, label = sample['image'], sample['label']
        # 随机旋转与翻转（一半概率）
        if random.random() > 0.5:
            image, label = random_rot_flip(image, label)
        elif random.random() > 0.5:
            image, label = random_rotate(image, label)
        x, y = image.shape[0], image.shape[1]
        # 缩放图片至输出尺寸
        if x != self.output_size[0] or y != self.output_size[1]:
            image = zoom(image, (self.output_size[0] / x, self.output_size[1] / y), order=3)  # why not 3?
            label = zoom(label, (self.output_size[0] / x, self.output_size[1] / y), order=0)
        image = torch.from_numpy(image.astype(np.float32)).unsqueeze(0)
        label = torch.from_numpy(label.astype(np.float32))
        sample = {'image': image, 'label': label.long()}
        return sample


class BraTs_dataset(Dataset):
    def __init__(self, base_dir, list_dir, split, transform=None):
        self.transform = transform  # using transform in torch!
        self.split = split
        self.sample_list = open(os.path.join(list_dir, self.split+'.txt')).readlines()
        self.data_dir = base_dir

    def __len__(self):
        return len(self.sample_list)

    # def __getitem__(self, idx):
    #     if self.split == "train":
    #         slice_name = self.sample_list[idx].strip('\n')
    #         image_data_path = os.path.join(self.data_dir, 'trainImage', slice_name+'.npy')
    #         image = np.load(image_data_path)
    #         image1 = image[:, :, 0]
    #         image2 = image[:, :, 1]
    #         image3 = image[:, :, 2]
    #         image4 = image[:, :, 3]
    #         cv2.imwrite('./test1.png', image1)
    #         cv2.imwrite('./test2.png', image2)
    #         cv2.imwrite('./test3.png', image3)
    #         cv2.imwrite('./test4.png', image4)
    #         label_data_path = os.path.join(self.data_dir, 'trainMask', slice_name+'.npy')
    #         label = np.load(label_data_path)
    #     else:
    #         vol_name = self.sample_list[idx].strip('\n')
    #         image_data_path = os.path.join(self.data_dir, 'testImage', vol_name + '.npy')
    #         image = np.load(image_data_path)
    #         label_data_path = os.path.join(self.data_dir, 'testMask', vol_name + '.npy')
    #         label = np.load(label_data_path)
    #
    #     sample = {'image': image, 'label': label}
    #     if self.transform:
    #         sample = self.transform(sample)
    #     sample['case_name'] = self.sample_list[idx].strip('\n')
    #     return sample
    def __getitem__(self, idx):
        slice_name = self.sample_list[idx].strip('\n')
        img_path = os.path.join(self.data_dir, self.split + 'Image', slice_name+'.npy')
        # img_path = self.img_paths[idx]
        mask_path = os.path.join(self.data_dir, self.split + 'Mask', slice_name+'.npy')
        # mask_path = self.mask_paths[idx]
        #读numpy数据(npy)的代码
        npimage = np.load(img_path)
        npmask = np.load(mask_path)
        npimage = npimage.transpose((2, 0, 1))

        WT_Label = npmask.copy()
        WT_Label[npmask == 1] = 1.
        WT_Label[npmask == 2] = 1.
        WT_Label[npmask == 4] = 1.
        TC_Label = npmask.copy()
        TC_Label[npmask == 1] = 1.
        TC_Label[npmask == 2] = 0.
        TC_Label[npmask == 4] = 1.
        ET_Label = npmask.copy()
        ET_Label[npmask == 1] = 0.
        ET_Label[npmask == 2] = 0.
        ET_Label[npmask == 4] = 1.
        nplabel = np.empty((160, 160, 3))
        nplabel[:, :, 0] = WT_Label
        nplabel[:, :, 1] = TC_Label
        nplabel[:, :, 2] = ET_Label
        nplabel = nplabel.transpose((2, 0, 1))

        nplabel = nplabel.astype("float32")
        npimage = npimage.astype("float32")

        sample = {'image': npimage[0], 'label': nplabel[0]}
        if self.transform:
            sample = self.transform(sample)
        sample['case_name'] = self.sample_list[idx].strip('\n')
        return sample
